// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compression/compression_internal.h>

#include <cstddef>
#include <cstdint>
#include <memory>

#include <brotli/decode.h>
#include <brotli/encode.h>
#include <brotli/types.h>


#include <turbo/utility/status.h>
#include <turbo/log/logging.h>
#include <turbo/base/macros.h>

namespace nebula::compress_internal {

    namespace {

        class BrotliDecompressor : public Decompressor {
        public:
            ~BrotliDecompressor() override {
                if (state_ != nullptr) {
                    BrotliDecoderDestroyInstance(state_);
                }
            }

            turbo::Status init() {
                state_ = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr);
                if (state_ == nullptr) {
                    return BrotliError("Brotli init failed");
                }
                return turbo::OkStatus();
            }

            turbo::Status reset() override {
                if (state_ != nullptr) {
                    BrotliDecoderDestroyInstance(state_);
                }
                return init();
            }

            turbo::Result<DecompressResult> Decompress(int64_t input_len, const uint8_t *input,
                                                       int64_t output_len, uint8_t *output) override {
                auto avail_in = static_cast<size_t>(input_len);
                auto avail_out = static_cast<size_t>(output_len);
                BrotliDecoderResult ret;

                ret = BrotliDecoderDecompressStream(state_, &avail_in, &input, &avail_out, &output,
                                                    nullptr /* total_out */);
                if (ret == BROTLI_DECODER_RESULT_ERROR) {
                    return BrotliError(BrotliDecoderGetErrorCode(state_), "Brotli decompress failed: ");
                }
                return DecompressResult{static_cast<int64_t>(input_len - avail_in),
                                        static_cast<int64_t>(output_len - avail_out),
                                        (ret == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT)};
            }

            bool IsFinished() override { return BrotliDecoderIsFinished(state_); }

        protected:
            turbo::Status BrotliError(const char *msg) { return turbo::io_error(msg); }

            turbo::Status BrotliError(BrotliDecoderErrorCode code, const char *prefix_msg) {
                return turbo::io_error(prefix_msg, BrotliDecoderErrorString(code));
            }

            BrotliDecoderState *state_ = nullptr;
        };

        // ----------------------------------------------------------------------
        // Brotli compressor implementation

        class BrotliCompressor : public Compressor {
        public:
            explicit BrotliCompressor(int compression_level, int window_bits)
                    : compression_level_(compression_level), window_bits_(window_bits) {}

            ~BrotliCompressor() override {
                if (state_ != nullptr) {
                    BrotliEncoderDestroyInstance(state_);
                }
            }

            turbo::Status init() {
                state_ = BrotliEncoderCreateInstance(nullptr, nullptr, nullptr);
                if (state_ == nullptr) {
                    return BrotliError("Brotli init failed");
                }
                if (!BrotliEncoderSetParameter(state_, BROTLI_PARAM_QUALITY, compression_level_)) {
                    return BrotliError("Brotli set compression level failed");
                }
                if (!BrotliEncoderSetParameter(state_, BROTLI_PARAM_LGWIN, window_bits_)) {
                    return BrotliError("Brotli set window size failed");
                }
                return turbo::OkStatus();
            }

            turbo::Result<CompressResult> Compress(int64_t input_len, const uint8_t *input,
                                                   int64_t output_len, uint8_t *output) override {
                auto avail_in = static_cast<size_t>(input_len);
                auto avail_out = static_cast<size_t>(output_len);
                BROTLI_BOOL ret;

                ret = BrotliEncoderCompressStream(state_, BROTLI_OPERATION_PROCESS, &avail_in, &input,
                                                  &avail_out, &output, nullptr /* total_out */);
                if (!ret) {
                    return BrotliError("Brotli compress failed");
                }
                return CompressResult{static_cast<int64_t>(input_len - avail_in),
                                      static_cast<int64_t>(output_len - avail_out)};
            }

            turbo::Result<FlushResult> flush(int64_t output_len, uint8_t *output) override {
                size_t avail_in = 0;
                const uint8_t *next_in = nullptr;
                auto avail_out = static_cast<size_t>(output_len);
                BROTLI_BOOL ret;

                ret = BrotliEncoderCompressStream(state_, BROTLI_OPERATION_FLUSH, &avail_in, &next_in,
                                                  &avail_out, &output, nullptr /* total_out */);
                if (!ret) {
                    return BrotliError("Brotli flush failed");
                }
                return FlushResult{static_cast<int64_t>(output_len - avail_out),
                                   !!BrotliEncoderHasMoreOutput(state_)};
            }

            turbo::Result<EndResult> End(int64_t output_len, uint8_t *output) override {
                size_t avail_in = 0;
                const uint8_t *next_in = nullptr;
                auto avail_out = static_cast<size_t>(output_len);
                BROTLI_BOOL ret;

                ret =
                        BrotliEncoderCompressStream(state_, BROTLI_OPERATION_FINISH, &avail_in, &next_in,
                                                    &avail_out, &output, nullptr /* total_out */);
                if (!ret) {
                    return BrotliError("Brotli end failed");
                }
                bool should_retry = !!BrotliEncoderHasMoreOutput(state_);
                        DKCHECK_EQ(should_retry, !BrotliEncoderIsFinished(state_));
                return EndResult{static_cast<int64_t>(output_len - avail_out), should_retry};
            }

        protected:
            turbo::Status BrotliError(const char *msg) { return turbo::io_error(msg); }

            BrotliEncoderState *state_ = nullptr;

        private:
            const int compression_level_;
            const int window_bits_;
        };

        // ----------------------------------------------------------------------
        // Brotli codec implementation

        class BrotliCodec : public Codec {
        public:
            explicit BrotliCodec(int compression_level, int window_bits)
                    : compression_level_(compression_level == kUseDefaultCompressionLevel
                                         ? kBrotliDefaultCompressionLevel
                                         : compression_level),
                      window_bits_(window_bits) {}

            turbo::Result<int64_t> Decompress(int64_t input_len, const uint8_t *input,
                                              int64_t output_buffer_len, uint8_t *output_buffer) override {
                        DKCHECK_GE(input_len, 0);
                        DKCHECK_GE(output_buffer_len, 0);
                std::size_t output_size = static_cast<size_t>(output_buffer_len);
                if (BrotliDecoderDecompress(static_cast<size_t>(input_len), input, &output_size,
                                            output_buffer) != BROTLI_DECODER_RESULT_SUCCESS) {
                    return turbo::io_error("Corrupt brotli compressed data.");
                }
                return output_size;
            }

            int64_t MaxCompressedLen(int64_t input_len,
                                     const uint8_t * TURBO_ARG_UNUSED(input)) override {
                        DKCHECK_GE(input_len, 0);
                return BrotliEncoderMaxCompressedSize(static_cast<size_t>(input_len));
            }

            turbo::Result<int64_t> Compress(int64_t input_len, const uint8_t *input,
                                            int64_t output_buffer_len, uint8_t *output_buffer) override {
                        DKCHECK_GE(input_len, 0);
                        DKCHECK_GE(output_buffer_len, 0);
                std::size_t output_size = static_cast<size_t>(output_buffer_len);
                if (BrotliEncoderCompress(compression_level_, window_bits_, BROTLI_DEFAULT_MODE,
                                          static_cast<size_t>(input_len), input, &output_size,
                                          output_buffer) == BROTLI_FALSE) {
                    return turbo::io_error("Brotli compression failure.");
                }
                return output_size;
            }

            turbo::Result<std::shared_ptr<Compressor>> MakeCompressor() override {
                auto ptr = std::make_shared<BrotliCompressor>(compression_level_, window_bits_);
                TURBO_RETURN_NOT_OK(ptr->init());
                return ptr;
            }

            turbo::Result<std::shared_ptr<Decompressor>> MakeDecompressor() override {
                auto ptr = std::make_shared<BrotliDecompressor>();
                TURBO_RETURN_NOT_OK(ptr->init());
                return ptr;
            }

            turbo::Status init() override {
                if (window_bits_ < BROTLI_MIN_WINDOW_BITS || window_bits_ > BROTLI_MAX_WINDOW_BITS) {
                    return turbo::invalid_argument_error("Brotli window_bits should be between ",
                                                  BROTLI_MIN_WINDOW_BITS, " and ", BROTLI_MAX_WINDOW_BITS);
                }
                return turbo::OkStatus();
            }

            CompressionType compression_type() const override { return CompressionType::BROTLI; }

            int compression_level() const override { return compression_level_; }

            int minimum_compression_level() const override { return BROTLI_MIN_QUALITY; }

            int maximum_compression_level() const override { return BROTLI_MAX_QUALITY; }

            int default_compression_level() const override {
                return kBrotliDefaultCompressionLevel;
            }

        private:
            const int compression_level_;
            const int window_bits_;
        };

    }  // namespace

    std::unique_ptr<Codec> MakeBrotliCodec(int compression_level,
                                           std::optional<int> window_bits) {
        return std::make_unique<BrotliCodec>(compression_level,
                                             window_bits.value_or(BROTLI_DEFAULT_WINDOW));
    }

}  // namespace nebula::compress_internal
